# -*- coding: utf-8 -*-
import asyncio
import os
import semantic_kernel as sk
from dotenv import load_dotenv, find_dotenv
from semantic_kernel.connectors.ai.open_ai.services.open_ai_chat_completion import OpenAIChatCompletion
from semantic_kernel.contents import ChatHistory
from semantic_kernel.functions import KernelArguments
from semantic_kernel.prompt_template import PromptTemplateConfig, InputVariable

_ = load_dotenv(find_dotenv())
# 创建 semantic kernel
kernel = sk.Kernel()
api_key = os.getenv("OPENAI_API_KEY")
serviceId = "default"
# 将 LLM 服务添加到 kernel 中
kernel.add_service(
    OpenAIChatCompletion(
        service_id=serviceId,
        ai_model_id="gpt-3.5-turbo-1106",
        api_key=api_key
    ),
)
# 例如我们要维护一个多轮对话，通过 request 和 history 两个变量分别存储 当前输入 和 对话历史
req_settings = kernel.get_service(serviceId).get_prompt_execution_settings_class()(service_id=serviceId)
# 提示语
prompt = """
{{$history}}
---
User:{{$request}}
Assistant:
"""
# 定义 Prompt 模板
# 模板中，变量以 {{$变量名}} 表示
prompt_template_config = PromptTemplateConfig(
    template=prompt,
    description="Multi-turn dialogue",
    execution_settings={serviceId: req_settings},
    input_variables=[
        InputVariable(name="request", description="The user input", is_required=True),
        InputVariable(name="history", description="The dialogue history", is_required=True),
    ],
)
# 注册 function
chat = kernel.add_function(
    function_name="chat",
    plugin_name="MyDemoPlugin",
    prompt_template_config=prompt_template_config,
)
chat_history = ChatHistory()
# 先添加一个系统消息
chat_history.add_system_message("You are a helpful chatbot who is good at answering user's questions.")


async def runAsyncFunc(*args):
    return await kernel.invoke(*args)


while True:
    request = input("User >")
    if not request:
        break
    result = asyncio.run(
        runAsyncFunc(
            chat,
            KernelArguments(
                request=request,
                history=chat_history
            ),
        )
    )
    print(f"Assistant > {result}")
    chat_history.add_user_message(request)
    chat_history.add_assistant_message(str(request))
